#!/usr/bin/env python
# coding: utf-8
import random
import math
import pandas as pd
import numpy as np
import csv
from time import time

# print函数输出所有，中间不会省略
# np.set_printoptions(threshold=np.inf)

start_time = time()


def write_list(out_list, mode="a"):
    """
    保存数据到CSV文件
    :param out_list:文件名字
    :param mode: 写入模式，默认是"a"，作用是从最后一行增加
    """
    # 保存到csv文件中
    with open("rating.csv", mode=mode, encoding="utf-8", newline="") as f:
        # 获取writer
        writer = csv.writer(f)
        # 写入第1行
        writer.writerow(out_list)


class RSVD():
    def __init__(self, allfile, trainfile, testfile, latentFactorNum=20, alpha_u=0.01, alpha_v=0.01, beta_u=0.01,
                 beta_v=0.01, learning_rate=0.01):
        """
        :param allfile:
            u.data，总数据
        :param trainfile:
            u1.base，训练集
        :param testfile:
            u1.train，测试集
        :param latentFactorNum:
            因子数量
        :param alpha_u:
             alpha_u的初始化
        :param alpha_v:
            alpha_v的初始化
        :param beta_u:
            beta_u的初始化
        :param beta_v:
            beta_vu的初始化
        :param learning_rate:
            训练速度
        """
        data_fields = ['user_id', 'item_id', 'rating', 'timestamp']
        # all data file
        allData = pd.read_table(allfile, names=data_fields)
        self.fullData = pd.read_table(allfile, names=data_fields)
        # training set file
        self.train_df = pd.read_table(trainfile, names=data_fields)
        # testing set file
        self.test_df = pd.read_table(testfile, names=data_fields)
        # get factor number
        self.latentFactorNum = latentFactorNum
        # get user number
        self.userNum = len(set(allData['user_id'].values))
        # get item number
        self.itemNum = len(set(allData['item_id'].values))
        # learning rate
        self.learningRate = learning_rate
        # the regularization lambda
        self.alpha_u = alpha_u
        self.alpha_v = alpha_v
        self.beta_u = beta_u
        self.beta_v = beta_v
        # initialize the model and parameters
        self.initModel()

    # initialize all parameters
    def initModel(self):
        """
        初始化模型参数u，bi，bu
        初始化方式为：
        u = 所有评分的平均数
        """
        print("initModel is start now")
        self.mu = self.train_df['rating'].mean()
        self.bu = np.zeros(self.userNum)
        self.bi = np.zeros(self.itemNum)
        self.U = np.mat(np.random.rand(self.userNum, self.latentFactorNum))
        self.V = np.mat(np.random.rand(self.itemNum, self.latentFactorNum))
        print("Initialize end.The user number is:%d,item number is:%d" % (self.userNum, self.itemNum))
        print("用户总数为:%d,项目总数为:%d" % (self.userNum, self.itemNum))

        print("initModel is end now")

    def train(self, iterTimes=100):
        """
        训练函数，训练过程中迭代更新u，bu，bi，U，V
        :param iterTimes: 迭代次数
        """
        print("train is start now")
        print("Beginning to train the model......")
        preRmse = 10000.0
        for iter in range(iterTimes):
            # print("self.bu = {}".format(self.bu[:10]))
            # print("self.bi = {}".format(self.bi[:10]))
            for index in self.train_df.index:
                if index % 20000 == 0:
                    print("第%s轮进度：%s%%" % (iter + 1, index / len(self.train_df.index) * 100))
                user = int(self.train_df.loc[index]['user_id']) - 1
                item = int(self.train_df.loc[index]['item_id']) - 1
                rating = float(self.train_df.loc[index]['rating'])
                pscore = self.predictScore(self.mu, self.bu[user], self.bi[item], self.U[user], self.V[item])
                eui = rating - pscore
                # update parameters bu and bi(user rating bais and item rating bais)
                # 更新 u，bu，bi
                self.mu = -eui#self.mu = pscore
                self.bu[user] += self.learningRate * (eui - self.beta_u * self.bu[user])
                self.bi[item] += self.learningRate * (eui - self.beta_v * self.bi[item])
                # 更新矩阵U，V
                temp = self.U[user]
                self.U[user] += self.learningRate * (eui * self.V[user] - self.alpha_u * self.U[user])
                self.V[item] += self.learningRate * (temp * eui - self.alpha_v * self.V[item])

            #                 print("self.bu = {}".format(self.bu.shape))
            #                 print("self.bi = {}".format(self.bi.shape))

            # calculate the current rmse
            curRmse = self.test(self.mu, self.bu, self.bi, self.U, self.V)
            print("当前迭代次数为：%d ,RMSE 为 : %f" % (iter + 1, curRmse))
            print("")
            if curRmse > preRmse:
                continue
            else:
                preRmse = curRmse
        print("Iteration finished!")
        print("train is end now")

    # test on the test set and calculate the RMSE
    def test(self, mu, bu, bi, U, V):
        """
        测试函数，每次迭代结束后都会进行测试
        使用测试集，计算 RMSE，最后返回RMSE
        :return: 测试集的 RMSE
        """
        print("test is start now")

        # 写入首行
        header = ["userid", "movieid", "rating","bu_add_bi_rate","ora_rating"]###########
        write_list(header, mode="w")

        # 初始化RMSE，计算预测分数
        cnt = self.test_df.shape[0]
        rmse = 0.0
        buT = bu.reshape(bu.shape[0], 1)
        predict_rate_matrix = mu + np.tile(buT, (1, self.itemNum)) + np.tile(bi, (self.userNum, 1)) + self.U * self.V.T

        # 开始对测试集计算RMSE
        print("开始对测试集计算RMSE")
        for i in self.test_df.index:
            user = int(self.test_df.loc[i]['user_id']) - 1
            item = int(self.test_df.loc[i]['item_id']) - 1
            score = float(self.test_df.loc[i]['rating'])
            pscore = predict_rate_matrix[user, item]
            rmse += math.pow(score - pscore, 2)
            # print("当前 userID 为：{}，movieID 为：{}，此时的实际分数为：{}，预测分数为：{}".format(user + 1, item + 1, score, pscore))

        # 计算最后的RMSE
        RMSE = math.sqrt(rmse / cnt)
        print("本次计算RMSE结束")

        # 开始写入数据到CSV格式文件中
        print("开始写入数据到CSV格式文件中")
        #predict_rate_matrix = mu + np.tile(buT, (1, self.itemNum)) + np.tile(bi, (self.userNum, 1))
        for i in self.fullData.index:
            user = int(self.fullData.loc[i]['user_id']) - 1
            item = int(self.fullData.loc[i]['item_id']) - 1
            score = float(self.fullData.loc[i]['rating'])
            pscore = predict_rate_matrix[user, item]
            if pscore < 1:
                pscore = 1
            if pscore > 5:
                pscore = 5
            rmse += math.pow(score - pscore, 2)
            predict_rate_matrix_bubi = np.tile(buT, (1, self.itemNum)) + np.tile(bi, (self.userNum, 1))
            bubi = predict_rate_matrix_bubi[user, item]
            row = [user + 1, item + 1, pscore, bubi, score]
            write_list(row)
        print("写入结束")

        print("test is end now")
        return RMSE

    def predictScore(self, mu, bu, bi, U, V):
        """
        训练过程中使用，作用是预测当前用户a对电影b的评分
        :return:预测分数，范围是1-5，包含小数
        """
        pscore = mu + bu + bi + np.multiply(U, V).sum()
        if pscore < 1:
            pscore = 1
        if pscore > 5:
            pscore = 5
        return pscore


if __name__ == '__main__':
    s = RSVD("datasets\\ml-100k\\u.data", "datasets\\ml-100k\\u1.base", "datasets\\ml-100k\\u1.test")
    s.train(iterTimes=10)

    end_time = time()
    print("总耗时时间为：{}秒".format(end_time - start_time))
